from flask import Flask, request, jsonify, send_file, make_response
from flask_cors import CORS
import io
import sys
import os
from PIL import Image
import logging
import base64
import torch
from torchvision import transforms

# 配置日志
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 导入transfer.py中的必要函数
means = [0.485, 0.456, 0.406]
stds = [0.229, 0.224, 0.225]
img_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(means, stds)
])

t_stds = torch.tensor(stds).cuda().half()[:,None,None]
t_means = torch.tensor(means).cuda().half()[:,None,None]

def tensor2im(var):
    return var.mul(t_stds).add(t_means).mul(255.).clamp(0, 255).permute(1, 2, 0)

def proc_pil_img(input_image, model):
    transformed_image = img_transforms(input_image)[None, ...].cuda().half()
    with torch.no_grad():
        result_image = model(transformed_image)[0]
        output_image = tensor2im(result_image)
        output_image = output_image.detach().cpu().numpy().astype('uint8')
        output_image = Image.fromarray(output_image)
    return output_image

app = Flask(__name__)

# 配置CORS
CORS(app, resources={
    r"/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization", "Accept", "Origin"]
    }
})

# 配置上传限制
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB 限制

# 加载模型 - 修改模型路径为相对于项目根目录的路径
model_path = "../models/ArcaneGANv0.4.jit"  # 使用一个简单的英文路径
logger.info(f"Loading model from: {model_path}")
model = torch.jit.load(model_path).eval().cuda().half()

def validate_image(file):
    """验证上传的图片文件"""
    allowed_extensions = {'.png', '.jpg', '.jpeg', '.webp'}
    ext = os.path.splitext(file.filename)[1].lower()
    if ext not in allowed_extensions:
        raise ValueError(f"不支持的文件类型: {ext}")
    
    image_bytes = file.read()
    file.seek(0)
    
    if len(image_bytes) > 10 * 1024 * 1024:  # 10MB
        raise ValueError("文件太大")
        
    try:
        image = Image.open(io.BytesIO(image_bytes))
        image.verify()
        return True
    except Exception as e:
        raise ValueError(f"无效的图像文件: {str(e)}")

@app.route('/upload', methods=['POST', 'OPTIONS'])
def upload_file():
    if request.method == 'OPTIONS':
        response = make_response()
        response.headers.add('Access-Control-Allow-Origin', '*')
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
        response.headers.add('Access-Control-Allow-Methods', 'POST')
        return response

    try:
        if 'file' not in request.files:
            return jsonify({"error": "没有上传文件"}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({"error": "未选择文件"}), 400
            
        try:
            validate_image(file)
        except ValueError as e:
            return jsonify({"error": str(e)}), 400
            
        # 读取和处理图像
        input_image = Image.open(file).convert('RGB')
        logger.info("开始处理图片")
        
        # 使用transfer.py中的处理函数
        stylized_image = proc_pil_img(input_image, model)
        
        # 转换为bytes
        img_byte_arr = io.BytesIO()
        stylized_image.save(img_byte_arr, format='JPEG')
        img_byte_arr = img_byte_arr.getvalue()
        
        logger.info("图片处理完成")

        # 创建响应
        response = make_response(send_file(
            io.BytesIO(img_byte_arr),
            mimetype='image/jpeg',
            as_attachment=True,
            download_name='cartoon_style.jpg'
        ))
        
        response.headers['Content-Type'] = 'image/jpeg'
        response.headers['Content-Length'] = len(img_byte_arr)
        response.headers['Cache-Control'] = 'no-cache'
        response.headers['Access-Control-Allow-Origin'] = '*'
        return response

    except Exception as e:
        logger.error(f"处理过程中出错: {str(e)}", exc_info=True)
        return jsonify({"error": str(e)}), 500

@app.route('/api/uploadimg', methods=['POST', 'OPTIONS'])
def upload_file_api():
    try:
        if 'file' not in request.files:
            return jsonify({"success": False, "error": "没有上传文件"}), 400
        file = request.files['file']
        if file.filename == '':
            return jsonify({"success": False, "error": "未选择文件"}), 400
            
        try:
            validate_image(file)
        except ValueError as e:
            return jsonify({"success": False, "error": str(e)}), 400
            
        # 读取和处理图像
        input_image = Image.open(file).convert('RGB')
        logger.info("开始处理图片")
        
        # 使用transfer.py中的处理函数
        stylized_image = proc_pil_img(input_image, model)
        
        # 转换为base64
        img_byte_arr = io.BytesIO()
        stylized_image.save(img_byte_arr, format='JPEG')
        img_byte_arr = img_byte_arr.getvalue()
        image_base64 = base64.b64encode(img_byte_arr).decode('utf-8')
        
        return jsonify({
            "success": True,
            "resultImagePath": f"data:image/jpeg;base64,{image_base64}"
        })

    except Exception as e:
        logger.error(f"处理过程中出错: {str(e)}", exc_info=True)
        return jsonify({
            "success": False,
            "error": str(e)
        }), 500

@app.route('/api/health', methods=['GET'])
def health_check():
    return jsonify({"status": "ok"}), 200

@app.errorhandler(413)
def request_entity_too_large(error):
    return jsonify({"error": "文件大小超过限制 (最大 16MB)"}), 413

if __name__ == '__main__':
    logger.info("Starting server...")
    app.run(debug=True, host='0.0.0.0', port=3001)
